﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Sigil.Impl
{
	internal class VerifiableTracker
	{
		public Int32 Iteration { get { return Transitions.Count; } }

		public Label BeganAt { get; private set; }

		// When the stack is "unbased" or "baseless", underflowing it results in wildcards
		//   eventually they'll be fixed up to actual types
		public Boolean IsBaseless { get; private set; }

		private List<InstructionAndTransitions> Transitions = new List<InstructionAndTransitions>();

		private Dictionary<Label, Int32> MarkedLabelsAtTransitions = new Dictionary<Label, Int32>();
		private Dictionary<Label, Int32> BranchesAtTransitions = new Dictionary<Label, Int32>();

		private Stack<List<TypeOnStack>> StartingStack = new Stack<List<TypeOnStack>>();

		public Boolean CanBePruned { get; private set; }

		public VerifiableTracker(Label beganAt, Boolean baseless = false, VerifiableTracker createdFrom = null, Boolean canBePruned = true)
		{
			IsBaseless = baseless;
			BeganAt = beganAt;
			CanBePruned = canBePruned;

			MarkedLabelsAtTransitions[beganAt] = 0;

			if (createdFrom != null)
			{
				StartingStack = GetStack(createdFrom);
			}
		}

		internal VerifiableTracker Concat(VerifiableTracker other)
		{
			var branchTo = BranchesAtTransitions.ContainsKey(other.BeganAt) ? BranchesAtTransitions[other.BeganAt] : Transitions.Count;
			var shouldTake = branchTo != Transitions.Count;

			var trans = new List<InstructionAndTransitions>(branchTo + other.Transitions.Count);

			if (shouldTake)
			{
				for (var i = 0; i < branchTo; i++)
				{
					trans.Add(Transitions[i]);
				}
			}
			else
			{
				trans.AddRange(Transitions);
			}

			trans.AddRange(other.Transitions);

			var canReuseCache = branchTo == Transitions.Count && IsBaseless && CachedVerifyStack != null;

			var ret =
					new VerifiableTracker(BeganAt, IsBaseless)
					{
						StartingStack = new Stack<List<TypeOnStack>>(StartingStack.Reverse()),
						Transitions = trans,
						CachedVerifyStack = canReuseCache ? new Stack<List<TypeOnStack>>(CachedVerifyStack.Reverse()) : null,
						CachedVerifyIndex = canReuseCache ? CachedVerifyIndex : null
					};

			return ret;
		}

		public VerificationResult Transition(InstructionAndTransitions legalTransitions)
		{
			Transitions.Add(legalTransitions);
			var ret = CollapseAndVerify();

			// revert!
			if (!ret.Success)
			{
				Transitions.RemoveAt(Transitions.Count - 1);
			}

			return ret;
		}

		public Int32? GetInstructionIndex(Int32 ix)
		{
			if (ix < 0 || ix >= Transitions.Count) throw new Exception("ix must be between 0 and " + (Transitions.Count - 1) + "; found " + ix);

			return Transitions[ix].InstructionIndex;
		}

		private static Stack<List<TypeOnStack>> GetStack(VerifiableTracker tracker)
		{
			var retStack = new Stack<List<TypeOnStack>>(tracker.StartingStack.Reverse());

			foreach (var t in tracker.Transitions)
			{
				UpdateStack(retStack, t, tracker.IsBaseless);
			}

			return retStack;
		}

		private static void UpdateStack(Stack<List<TypeOnStack>> stack, InstructionAndTransitions wrapped, Boolean isBaseless)
		{
			var legal = wrapped.Transitions;
			var instr = wrapped.Instruction;

			var legalSize = 0;

			legal.ForEach(
					t =>
					{
						legalSize += t.PushedToStack.Length;

						if (t.Before != null) t.Before(stack, isBaseless);
					}
			);

			if (legal.Any(l => l.PoppedFromStack.Any(u => u == TypeOnStack.Get<PopAllType>())))
			{
				if (instr.HasValue)
				{
					for (var i = 0; i < stack.Count; i++)
					{
						var ix = stack.Count - i - 1;
						stack.ElementAt(i).ForEach(y => y.Mark(wrapped, ix));
					}
				}

				stack.Clear();
			}
			else
			{
				var toPop = legal.First().PoppedCount;

				for (var j = 0; j < toPop && stack.Count > 0; j++)
				{
					var popped = stack.Pop();

					if (instr.HasValue)
					{
						var ix = toPop - j - 1;
						popped.ForEach(y => y.Mark(wrapped, ix));
					}
				}
			}

			var toPush = new List<TypeOnStack>(legalSize);
			var pushed = new HashSet<TypeOnStack>();
			for (var i = 0; i < legal.Count; i++)
			{
				foreach (var p in legal[i].PushedToStack)
				{
					if (pushed.Contains(p)) continue;

					toPush.Add(p);
					pushed.Add(p);
				}
			}

			if (toPush.Count > 0)
			{
				stack.Push(toPush);
			}
		}

		private List<StackTransition> GetLegalTransitions(List<StackTransition> ops, Stack<List<TypeOnStack>> runningStack)
		{
			var ret = new List<StackTransition>(ops.Count);

			for (var i = 0; i < ops.Count; i++)
			{
				var w = ops[i];

				if (w.PoppedFromStack.All(u => u == TypeOnStack.Get<PopAllType>()))
				{
					ret.Add(w);
					continue;
				}

				var onStack = runningStack.Peek(IsBaseless, w.PoppedCount);

				if (onStack == null)
				{
					continue;
				}

				if (w.PushedToStack.Any(p => p == TypeOnStack.Get<SamePointerType>()))
				{
					if (w.PushedToStack.Length > 1)
					{
						throw new Exception("SamePointerType can be only product of a transition which contains it");
					}

					var shouldBePointer = onStack.SelectMany(p => p.Where(x => x.IsPointer || x == TypeOnStack.Get<WildcardType>())).Distinct();

					if (!shouldBePointer.Any()) continue;
					w = new StackTransition(w.PoppedFromStack, new[] { shouldBePointer.Single() });
				}

				if (w.PushedToStack.Any(p => p == TypeOnStack.Get<SameByRefType>()))
				{
					if (w.PushedToStack.Length > 1)
					{
						throw new Exception("SameByRefType can be only product of a transition which contains it");
					}

					var shouldBeByRef = onStack.SelectMany(p => p.Where(x => x.IsReference || x == TypeOnStack.Get<WildcardType>())).Distinct();

					if (!shouldBeByRef.Any()) continue;
					w = new StackTransition(w.PoppedFromStack, new[] { shouldBeByRef.Single() });
				}

				Boolean outerContinue = false;

				for (var j = 0; j < w.PoppedCount; j++)
				{
					var shouldBe = w.PoppedFromStack[j];
					var actuallyIs = onStack[j];

					if (!actuallyIs.Any(a => shouldBe.IsAssignableFrom(a)))
					{
						outerContinue = true;
						break;
					}
				}

				if (outerContinue) continue;

				ret.Add(w);
			}

			return ret;
		}

		private Stack<List<TypeOnStack>> CachedVerifyStack;
		private Int32? CachedVerifyIndex;

		public VerificationResult CollapseAndVerify()
		{
			var runningStack = CachedVerifyStack ?? new Stack<List<TypeOnStack>>(StartingStack.Reverse());

			Int32 i = CachedVerifyIndex ?? 0;

			for (; i < Transitions.Count; i++)
			{
				var wrapped = Transitions[i];
				var ops = wrapped.Transitions;

				if (ops.Any(o => o.StackSizeMustBe.HasValue))
				{
					if (ops.Count > 1)
					{
						throw new Exception("Shouldn't have multiple 'must be size' transitions at the same point");
					}

					var doIt = ops[0];

					if (doIt.StackSizeMustBe != runningStack.Count)
					{
						return VerificationResult.FailureStackSize(this, i, doIt.StackSizeMustBe.Value);
					}
				}

				var legal = GetLegalTransitions(ops, runningStack);

				if (legal.Count == 0)
				{
					var wouldPop = ops.GroupBy(g => g.PoppedFromStack.Length).Single().Key;

					if (runningStack.Count < wouldPop)
					{
						return VerificationResult.FailureUnderflow(this, i, wouldPop, runningStack);
					}

					IEnumerable<TypeOnStack> expected;
					var stackI = FindStackFailureIndex(runningStack, ops, out expected);

					return VerificationResult.FailureTypeMismatch(this, i, stackI, expected, runningStack);
				}

				if (legal.GroupBy(g => new { a = g.PoppedCount, b = g.PushedToStack.Length }).Count() > 1)
				{
					throw new Exception("Shouldn't be possible; legal transitions should have same push/pop #s");
				}

				// No reason to do all this work again
				Transitions[i] = new InstructionAndTransitions(wrapped.Instruction, wrapped.InstructionIndex, legal);

				Boolean popAll = legal.Any(l => l.PoppedFromStack.Contains(TypeOnStack.Get<PopAllType>()));
				if (popAll && legal.Count() != 1)
				{
					throw new Exception("PopAll cannot coexist with any other transitions");
				}

				if (!popAll)
				{
					var toPop = legal.First().PoppedCount;

					if (toPop > runningStack.Count && !IsBaseless)
					{
						return VerificationResult.FailureUnderflow(this, i, toPop, runningStack);
					}
				}

				Boolean isDuplicate = legal.Any(l => l.IsDuplicate);
				if (isDuplicate && legal.Count() > 1)
				{
					throw new Exception("Duplicate must be only transition");
				}

				if (isDuplicate)
				{
					if (!IsBaseless && runningStack.Count == 0) return VerificationResult.FailureUnderflow(this, i, 1, runningStack);

					var toPush = runningStack.Count > 0 ? runningStack.Peek() : new List<TypeOnStack>(new[] { TypeOnStack.Get<WildcardType>() });

					UpdateStack(runningStack, new InstructionAndTransitions(wrapped.Instruction, wrapped.InstructionIndex, new List<StackTransition>(new[] { new StackTransition(new TypeOnStack[0], toPush) })), IsBaseless);
				}
				else
				{
					UpdateStack(runningStack, new InstructionAndTransitions(wrapped.Instruction, wrapped.InstructionIndex, legal), IsBaseless);
				}
			}

			CachedVerifyIndex = i;
			CachedVerifyStack = runningStack;

			return VerificationResult.Successful(this, runningStack);
		}

		private Int32 FindStackFailureIndex(Stack<List<TypeOnStack>> types, IEnumerable<StackTransition> ops, out IEnumerable<TypeOnStack> expected)
		{
			var stillLegal = new List<StackTransition>(ops);

			for (var i = 0; i < types.Count; i++)
			{
				var actuallyIs = types.ElementAt(i);

				var legal = stillLegal.Where(l => actuallyIs.Any(a => l.PoppedFromStack[i].IsAssignableFrom(a)));

				if (!legal.Any())
				{
					expected = stillLegal.Select(l => l.PoppedFromStack[i]).Distinct();
					return i;
				}

				stillLegal = new List<StackTransition>(legal);
			}

			throw new Exception("Shouldn't be possible");
		}

		public VerifiableTracker Clone()
		{
			return
					new VerifiableTracker(BeganAt)
					{
						IsBaseless = IsBaseless,
						MarkedLabelsAtTransitions = new Dictionary<Label, Int32>(MarkedLabelsAtTransitions),
						BranchesAtTransitions = new Dictionary<Label, Int32>(BranchesAtTransitions),
						Transitions = new List<InstructionAndTransitions>(Transitions),
						StartingStack = new Stack<List<TypeOnStack>>(StartingStack.Reverse())
					};
		}

		// Returns the current stack *if* it can be inferred down to single types *and* is either based or verifiable to the given depth
		public Stack<TypeOnStack> InferStack(Int32 ofDepth)
		{
			var res = CollapseAndVerify();

			if (res.Stack.Count < ofDepth) return null;

			var ret = new Stack<TypeOnStack>();
			for (var i = ofDepth - 1; i >= 0; i--)
			{
				var couldBe = res.Stack.ElementAt(i);

				if (couldBe.Count() > 1) return null;

				ret.Push(couldBe.Single());
			}

			return ret;
		}

		public override String ToString()
		{
			var ret = new StringBuilder();

			if (StartingStack.Count > 0)
			{
				ret.AppendLine(
						"starts with: " +
								String.Join(", ",
										StartingStack.Select(s => "[" + String.Join(", or", s.Select(x => x.ToString()).ToArray()) + "]").ToArray()
								)
				);
			}

			for (var i = 0; i < Transitions.Count; i++)
			{
				var label = MarkedLabelsAtTransitions.Where(kv => kv.Value == i).Select(kv => kv.Key).SingleOrDefault();

				if (label != null)
				{
					ret.AppendLine(label.Name + ":");
				}

				var tran = Transitions[i];

				ret.AppendLine(tran.ToString());
			}

			return ret.ToString();
		}
	}
}